package ws

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"net/http"
	"sync"
	"time"

	"websocket-binance/internal/config"

	"github.com/gorilla/websocket"
)

// WebSocketServer 实现高性能的WebSocket服务器

type WebSocketServer struct {
	config       *config.ServerConfig
	clients      map[string]*Client
	clientsMutex sync.Mutex
	upgrader     websocket.Upgrader
	ticker       *time.Ticker
	ctx          context.Context
	cancel       context.CancelFunc
}

// Client 表示一个WebSocket客户端连接

type Client struct {
	conn          *websocket.Conn
	subscriptions map[string]bool // 订阅的频道
	send          chan []byte
}

// Message 定义客户端消息格式
type Message struct {
	Type    string          `json:"type"` // subscribe, unsubscribe, ping, pong
	Channel string          `json:"channel,omitempty"`
	Data    json.RawMessage `json:"data,omitempty"`
}

// NewWebSocketServer 创建一个新的WebSocket服务器
func NewWebSocketServer(cfg *config.ServerConfig) *WebSocketServer {
	ctx, cancel := context.WithCancel(context.Background())

	// 解析超时设置，但不在此直接使用，因为它们会在HTTP服务器和连接处理中分别使用

	server := &WebSocketServer{
		config:  cfg,
		clients: make(map[string]*Client),
		upgrader: websocket.Upgrader{
			ReadBufferSize:  1024,
			WriteBufferSize: 1024,
			CheckOrigin: func(r *http.Request) bool {
				// 允许所有来源
				return true
			},
		},
		ctx:    ctx,
		cancel: cancel,
	}

	// 设置心跳定时器 - 心跳间隔设置为read_timeout的1/2，确保有足够时间收到pong响应
	readTimeout := parseDuration(server.config.ReadTimeout)
	heartbeatInterval := readTimeout / 2
	server.ticker = time.NewTicker(heartbeatInterval)
	log.Printf("设置心跳间隔为: %v，读取超时为: %v", heartbeatInterval, readTimeout)
	go server.heartbeatLoop()

	return server
}

// Start 启动WebSocket服务器
func (s *WebSocketServer) Start() error {
	http.HandleFunc("/ws", s.handleWebSocket)

	server := &http.Server{
		Addr:    fmt.Sprintf(":%d", s.config.Port),
		Handler: nil,
		// 解析并设置超时
		ReadTimeout:       parseDuration(s.config.ReadTimeout),
		WriteTimeout:      parseDuration(s.config.WriteTimeout),
		IdleTimeout:       parseDuration(s.config.IdleTimeout),
		MaxHeaderBytes:    1 << 20,
		TLSConfig:         nil,
		ReadHeaderTimeout: 0,
		ConnState:         nil,
		ErrorLog:          nil,
		BaseContext:       nil,
		ConnContext:       nil,
	}

	go func() {
		<-s.ctx.Done()
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		server.Shutdown(shutdownCtx)
	}()

	log.Printf("WebSocket服务器启动在端口: %d", s.config.Port)
	return server.ListenAndServe()
}

// handleWebSocket 处理WebSocket连接请求
func (s *WebSocketServer) handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := s.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket升级失败: %v", err)
		return
	}

	// 为客户端生成唯一ID
	clientID := generateClientID()

	// 创建客户端 - 增加发送缓冲区大小以减少因缓冲区满而关闭连接的可能性
	client := &Client{
		conn:          conn,
		subscriptions: make(map[string]bool),
		send:          make(chan []byte, 1024), // 增加缓冲区大小从256到1024
	}

	// 添加客户端到客户端列表
	s.clientsMutex.Lock()
	s.clients[clientID] = client
	s.clientsMutex.Unlock()

	log.Printf("新客户端连接: %s, 当前连接数: %d", clientID, len(s.clients))

	// 启动客户端的读写循环
	go s.writePump(clientID, client)
	go s.readPump(clientID, client)
}

// readPump 处理从客户端读取消息
func (s *WebSocketServer) readPump(clientID string, client *Client) {
	defer func() {
		s.unregisterClient(clientID)
		client.conn.Close()
	}()

	client.conn.SetReadLimit(int64(s.config.MaxMessageSize))

	// 设置读取超时
	readTimeout := parseDuration(s.config.ReadTimeout)
	client.conn.SetReadDeadline(time.Now().Add(readTimeout))
	client.conn.SetPongHandler(func(string) error {
		// 收到pong消息时更新读取超时
		client.conn.SetReadDeadline(time.Now().Add(readTimeout))
		log.Printf("客户端 %s 收到pong消息，延长读取超时至 %v 后", clientID, time.Now().Add(readTimeout).Format("15:04:05"))
		return nil
	})

	for {
		_, message, err := client.conn.ReadMessage()
		if err != nil {
			// 增加更详细的错误日志，帮助诊断断开原因
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				log.Printf("客户端 %s 读取超时断开连接: %v", clientID, err)
			} else if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("客户端 %s 意外关闭错误: %v", clientID, err)
			} else {
				log.Printf("客户端 %s 读取消息错误: %v", clientID, err)
			}
			break
		}

		// 处理客户端消息
		s.handleClientMessage(clientID, client, message)
	}
}

// writePump 处理向客户端发送消息
func (s *WebSocketServer) writePump(clientID string, client *Client) {
	defer func() {
		// 提供明确的关闭代码
		client.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "Server shutdown"))
		client.conn.Close()
	}()

	for {
		select {
		case message, ok := <-client.send:
			// 设置写入超时
			writeTimeout, _ := time.ParseDuration(s.config.WriteTimeout)
			client.conn.SetWriteDeadline(time.Now().Add(writeTimeout))
			if !ok {
				// 服务器关闭了channel
				client.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			w, err := client.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			w.Write(message)

			// 发送队列中的所有消息
			n := len(client.send)
			for i := 0; i < n; i++ {
				w.Write([]byte{'\n'})
				w.Write(<-client.send)
			}

			if err := w.Close(); err != nil {
				return
			}
		case <-s.ctx.Done():
			return
		}
	}
}

// handleClientMessage 处理客户端发送的消息
func (s *WebSocketServer) handleClientMessage(clientID string, client *Client, message []byte) {
	var msg Message
	if err := json.Unmarshal(message, &msg); err != nil {
		log.Printf("解析客户端消息失败: %v", err)
		return
	}

	switch msg.Type {
	case "subscribe":
		if msg.Channel != "" {
			client.subscriptions[msg.Channel] = true
			log.Printf("客户端 %s 订阅了频道: %s", clientID, msg.Channel)
		}
	case "unsubscribe":
		if msg.Channel != "" {
			delete(client.subscriptions, msg.Channel)
			log.Printf("客户端 %s 取消订阅了频道: %s", clientID, msg.Channel)
		}
	case "ping":
		// 响应pong
		pongMsg := Message{
			Type: "pong",
		}
		pongBytes, _ := json.Marshal(pongMsg)
		client.send <- pongBytes
	}
}

// Publish 发布消息到指定频道
func (s *WebSocketServer) Publish(channel string, data []byte) {
	s.clientsMutex.Lock()
	defer s.clientsMutex.Unlock()

	// 创建发布消息
	publishMsg := Message{
		Type:    "publish",
		Channel: channel,
		Data:    data,
	}

	msgBytes, err := json.Marshal(publishMsg)
	if err != nil {
		log.Printf("序列化发布消息失败: %v", err)
		return
	}

	// 发送消息给所有订阅了该频道的客户端
	for clientID, client := range s.clients {
		if client.subscriptions[channel] {
			select {
			case client.send <- msgBytes:
			default:
				// 如果客户端的发送缓冲区已满，先尝试发送关闭消息
				client.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseServiceRestart, "Buffer full"))
				client.conn.Close()
				close(client.send)
				// 已经在锁内，不需要再次加锁
				delete(s.clients, clientID)
			}
		}
	}
}

// heartbeatLoop 定期发送心跳消息给所有客户端
func (s *WebSocketServer) heartbeatLoop() {
	pingMsg := Message{
		Type: "ping",
	}
	pingBytes, _ := json.Marshal(pingMsg)
	readTimeout := parseDuration(s.config.ReadTimeout)

	for {
		select {
		case <-s.ticker.C:
			currentTime := time.Now()
			log.Printf("发送心跳消息给所有客户端，当前时间: %v", currentTime.Format("15:04:05"))
			s.clientsMutex.Lock()
			for clientID, client := range s.clients {
				// 手动设置连接的读取超时，确保与pong消息处理保持一致
				expectedPongTime := currentTime.Add(readTimeout).Format("15:04:05")
				client.conn.SetReadDeadline(currentTime.Add(readTimeout))
				
				select {
				case client.send <- pingBytes:
					log.Printf("已向客户端 %s 发送ping消息，期望pong响应超时时间: %v", clientID, expectedPongTime)
				default:
					// 如果客户端的发送缓冲区已满，先尝试发送关闭消息
					log.Printf("客户端 %s 发送缓冲区已满，关闭连接", clientID)
					client.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseServiceRestart, "Buffer full"))
					client.conn.Close()
					close(client.send)
					delete(s.clients, clientID)
				}
			}
			s.clientsMutex.Unlock()
		case <-s.ctx.Done():
			s.ticker.Stop()
			return
		}
	}
}

// unregisterClient 从客户端列表中移除客户端
func (s *WebSocketServer) unregisterClient(clientID string) {
	s.clientsMutex.Lock()
	defer s.clientsMutex.Unlock()

	if client, ok := s.clients[clientID]; ok {
		close(client.send)
		delete(s.clients, clientID)
		log.Printf("客户端断开连接: %s, 当前连接数: %d", clientID, len(s.clients))
	}
}

// Stop 停止WebSocket服务器
func (s *WebSocketServer) Stop() {
	s.cancel()
}

// 辅助函数：生成客户端ID
func generateClientID() string {
	return fmt.Sprintf("client-%d", time.Now().UnixNano())
}

// 辅助函数：解析持续时间字符串，出错时返回默认值
func parseDuration(durationStr string) time.Duration {
	duration, err := time.ParseDuration(durationStr)
	if err != nil {
		log.Printf("解析持续时间失败: %v, 使用默认值5秒", err)
		return 5 * time.Second
	}
	return duration
}
